#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO] $1${NC}"
}

log_warn() {
    echo -e "${YELLOW}[WARN] $1${NC}"
}

log_error() {
    echo -e "${RED}[ERROR] $1${NC}"
}

# 检查命令是否存在
check_command() {
    if ! command -v $1 &> /dev/null; then
        log_error "$1 命令未找到，请先安装"
        exit 1
    fi
}

# 检查 HDFS 服务状态
check_hdfs_status() {
    if ! hdfs dfsadmin -report &> /dev/null; then
        log_error "HDFS 服务未运行"
        exit 1
    fi
}

# 节点管理函数
add_datanode() {
    local node_name=$1
    local node_ip=$2
    local namenode_ip=$3

    log_info "开始添加 DataNode: $node_name ($node_ip)"

    # 检查节点连通性
    if ! ping -c 1 $node_ip &> /dev/null; then
        log_error "无法连接到节点 $node_ip"
        return 1
    fi

    # 在远程节点上安装 Hadoop
    ssh $node_ip "sudo apt update && sudo apt install -y hadoop"

    # 配置远程节点
    ssh $node_ip "sudo tee /etc/hadoop/core-site.xml << EOF
<?xml version=\"1.0\" encoding=\"UTF-8\"?>
<configuration>
    <property>
        <name>fs.defaultFS</name>
        <value>hdfs://$namenode_ip:9000</value>
    </property>
</configuration>
EOF"

    ssh $node_ip "sudo tee /etc/hadoop/hdfs-site.xml << EOF
<?xml version=\"1.0\" encoding=\"UTF-8\"?>
<configuration>
    <property>
        <name>dfs.datanode.data.dir</name>
        <value>/var/lib/hadoop-hdfs/data</value>
    </property>
    <property>
        <name>dfs.datanode.address</name>
        <value>$node_ip:50010</value>
    </property>
    <property>
        <name>dfs.datanode.http.address</name>
        <value>$node_ip:50075</value>
    </property>
</configuration>
EOF"

    # 启动远程节点服务
    ssh $node_ip "sudo systemctl daemon-reload && sudo systemctl enable hadoop-datanode && sudo systemctl start hadoop-datanode"

    # 验证节点加入
    sleep 10
    if hdfs dfsadmin -report | grep -q $node_ip; then
        log_info "DataNode $node_name 成功加入集群"
    else
        log_error "DataNode $node_name 加入集群失败"
        return 1
    fi
}

remove_datanode() {
    local node_name=$1
    local node_ip=$2

    log_info "开始移除 DataNode: $node_name ($node_ip)"

    # 检查节点是否存在
    if ! hdfs dfsadmin -report | grep -q $node_ip; then
        log_error "DataNode $node_name 不在集群中"
        return 1
    fi

    # 将节点加入退役列表
    echo $node_ip >> /etc/hadoop/dfs.exclude

    # 刷新节点列表
    hdfs dfsadmin -refreshNodes

    # 等待数据迁移完成
    sleep 30

    # 停止远程节点服务
    ssh $node_ip "sudo systemctl stop hadoop-datanode && sudo systemctl disable hadoop-datanode"

    # 清理远程节点数据
    ssh $node_ip "sudo rm -rf /var/lib/hadoop-hdfs/data/*"

    # 从退役列表中移除节点
    sed -i "/$node_ip/d" /etc/hadoop/dfs.exclude
    hdfs dfsadmin -refreshNodes

    log_info "DataNode $node_name 已成功移除"
}

# 磁盘管理函数
add_disk() {
    local node_name=$1
    local disk_path=$2

    log_info "开始为节点 $node_name 添加磁盘: $disk_path"

    # 检查节点是否存在
    if ! hdfs dfsadmin -report | grep -q $node_name; then
        log_error "节点 $node_name 不在集群中"
        return 1
    fi

    # 在节点上创建数据目录
    ssh $node_name "sudo mkdir -p $disk_path && sudo chown -R hdfs:hdfs $disk_path"

    # 更新 HDFS 配置
    ssh $node_name "sudo tee -a /etc/hadoop/hdfs-site.xml << EOF
    <property>
        <name>dfs.datanode.data.dir</name>
        <value>$disk_path</value>
    </property>
EOF"

    # 重启节点服务
    ssh $node_name "sudo systemctl restart hadoop-datanode"

    # 验证磁盘添加
    sleep 10
    if hdfs dfsadmin -report | grep -q $disk_path; then
        log_info "磁盘 $disk_path 已成功添加到节点 $node_name"
    else
        log_error "磁盘添加失败"
        return 1
    fi
}

handle_failed_disk() {
    local node_name=$1
    local disk_path=$2

    log_info "开始处理节点 $node_name 的故障磁盘: $disk_path"

    # 检查节点是否存在
    if ! hdfs dfsadmin -report | grep -q $node_name; then
        log_error "节点 $node_name 不在集群中"
        return 1
    fi

    # 从配置中移除故障磁盘
    ssh $node_name "sudo sed -i '/$disk_path/d' /etc/hadoop/hdfs-site.xml"

    # 重启节点服务
    ssh $node_name "sudo systemctl restart hadoop-datanode"

    # 等待数据迁移完成
    sleep 30

    # 卸载故障磁盘
    ssh $node_name "sudo umount $disk_path"

    # 标记磁盘为故障
    ssh $node_name "sudo echo 'DISK_FAILED' > $disk_path/disk_status"

    log_info "故障磁盘 $disk_path 已成功处理"
}

# 主函数
main() {
    # 检查必要命令
    check_command hdfs
    check_command ssh
    check_command systemctl

    # 检查 HDFS 状态
    check_hdfs_status

    case "$1" in
        "add-datanode")
            if [ $# -ne 4 ]; then
                log_error "用法: $0 add-datanode <node_name> <node_ip> <namenode_ip>"
                exit 1
            fi
            add_datanode "$2" "$3" "$4"
            ;;
        "remove-datanode")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 remove-datanode <node_name> <node_ip>"
                exit 1
            fi
            remove_datanode "$2" "$3"
            ;;
        "add-disk")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 add-disk <node_name> <disk_path>"
                exit 1
            fi
            add_disk "$2" "$3"
            ;;
        "handle-failed-disk")
            if [ $# -ne 3 ]; then
                log_error "用法: $0 handle-failed-disk <node_name> <disk_path>"
                exit 1
            fi
            handle_failed_disk "$2" "$3"
            ;;
        *)
            log_error "用法: $0 {add-datanode|remove-datanode|add-disk|handle-failed-disk}"
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@" 